import tqdm
from gxl_ai_utils.utils import utils_file
import numpy as np

def get_token_str_from_npy(npy_path):
    token_list = np.load(npy_path).tolist()
    token_str = " ".join([str(token) for token in token_list])
    return token_str

list_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/train_wtn.list"
dict_list = utils_file.load_dict_list_from_jsonl(list_path)
token_scp_dict = {}
def func_little(res_dict, input_dict_list):
    for input_dict in tqdm.tqdm(input_dict_list, total=len(input_dict_list)):
        key = input_dict["key"]
        npy_path = input_dict["npy"]
        token_str = get_token_str_from_npy(npy_path)
        res_dict[key] = token_str

runner = utils_file.GxlDynamicThreadPool()
input_dict_list_list = utils_file.do_split_list(dict_list, 20)
for input_dict_list in input_dict_list_list:
    runner.add_task(func_little, [token_scp_dict, input_dict_list])
runner.start()

token_scp_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/train_token.scp"
utils_file.write_dict_to_scp(token_scp_dict, token_scp_path)

txt_scp_dict = {}
for dict_item in tqdm.tqdm(dict_list):
    txt = dict_item["txt"]
    key = dict_item["key"]
    txt_scp_dict[key] = txt
token_scp_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/train_text.scp"
utils_file.write_dict_to_scp(txt_scp_dict, token_scp_path)

wcp_scp_dict = {}
for dict_item in tqdm.tqdm(dict_list):
    wav = dict_item["wav"]
    key = dict_item["key"]
    txt_scp_dict[key] = wav
token_scp_path = "/home/work_nfs8/xlgeng/new_workspace/wenet_gxl_salmonn_TTS/examples/aishell/wenetspeech4tts_handler/data_list/train_wav.scp"
utils_file.write_dict_to_scp(txt_scp_dict, token_scp_path)
